/* Copyright (C) Nial Peters 2013
 *
 * This file is part of pydoas.
 *
 * pydoas is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * pydoas is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with pydoas.  If not, see <http://www.gnu.org/licenses/>.
 *
 * The code in this file is modified from
 * http://blog.enthought.com/python/numpy-arrays-with-pre-allocated-memory/
 */

#include <Python.h>
#include "structmember.h"
#include <numpy/arrayobject.h>
#include "autodealloc_numpy.h"
#include <stdlib.h>
#include "levmar.h"
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_fit.h>
#include <gsl/gsl_blas.h>

/******************************************************************************/

typedef struct FitData{
	int nrows; //number of rows in fitting_params matrix
	int ncols; //number of columns in fitting_params matrix
	int spec_size; //size of measured spectrum
	gsl_matrix_view lin_fit_params;
	double *measured_spec; //full spectrum data (i.e. not clipped to fitting range)
	double *wavelengths; //wavelengths corresponding to measured_spec values
	double *fitting_wavelengths;
	double *sns_wavelengths; //workspace for storing shifted and squeezed values
	gsl_multifit_linear_workspace *workspace;
	gsl_vector *lin_fit_result;
	gsl_matrix *cov;
} fit_data_t;

/******************************************************************************/
static fit_data_t *
createFitDataObj(PyObject * A, PyObject *spec, PyObject *wavelengths,
				 PyObject *fitting_wavelengths){

	int spec_size = PyArray_DIM(spec, 0);

	fit_data_t *fd = malloc(sizeof(fit_data_t));
	fd->measured_spec = PyArray_DATA(spec);
	fd->wavelengths = PyArray_DATA(wavelengths);
	fd->fitting_wavelengths = PyArray_DATA(fitting_wavelengths);
	fd->sns_wavelengths = malloc(spec_size * sizeof(double));

	fd->ncols = PyArray_DIM(A, 1);
	fd->nrows = PyArray_DIM(A, 0);
	fd->spec_size = spec_size;

	//now create a matrix from the fitting data array
	fd->lin_fit_params = gsl_matrix_view_array (PyArray_DATA(A), fd->nrows, fd->ncols);
	fd->lin_fit_result = gsl_vector_alloc (fd->ncols);
	fd->cov = gsl_matrix_alloc(fd->ncols,fd->ncols);
	fd->workspace = gsl_multifit_linear_alloc(fd->nrows, fd->ncols);

	return fd;
}

/******************************************************************************/

static void
freeFitData(fit_data_t **fd){

	if ((fd == NULL) || (*fd == NULL)){
		return;
	}
	gsl_multifit_linear_free((*fd)->workspace);
	free((*fd)->sns_wavelengths);
	free(*fd);
	*fd = NULL;

}

/******************************************************************************/

static inline int
linearInterpolate(const double *x, const double *y,
				  const int n, const double *interp_x,
				  double *interp_y, const int interp_n)
{
	/*
	 * Linear interpolates a data series described by x and y (whose length
	 * is n), onto the points interp_x (which must be of length interp_n). The
	 * results are stored in interp_y. Returns 0 on success, 1 on failure.
	 *
	 * max(interp_x) <= max(x) and min(interp_x) >= min(x)
	 */

	register int x_idx=0;
	register int interp_idx;

	//error checking
	if ((x[0] > interp_x[0]) || (x[n-1] < interp_x[interp_n - 1])){
		return 1;
	}

	for (interp_idx=0; interp_idx<interp_n; interp_idx++){

		while (interp_x[interp_idx] > x[x_idx]){
			x_idx++;
		}
		interp_y[interp_idx] = y[x_idx-1] + (interp_x[interp_idx]-x[x_idx-1]) *
				(y[x_idx] - y[x_idx-1])/(x[x_idx] - x[x_idx-1]);
	}

	return 0;
}

/******************************************************************************/
static inline void
fit_func(double *p, double *x, int m, int n, void *data)
{
	/*
	 *
	 */
	register int i;
	fit_data_t *fd = (fit_data_t *)data;
	double chisq;

	//shift and squeeze the wavelength data
	for(i=0; i<fd->spec_size; i++){
		fd->sns_wavelengths[i] = p[0] +  (fd->wavelengths[i]);
	}

	//linear interpolate the spectrum to be at the shifted and squeezed points
	if(linearInterpolate(fd->sns_wavelengths, fd->measured_spec, fd->spec_size, fd->fitting_wavelengths, x, fd->nrows)){
		printf("Error!\n");
		exit(1);
		return;
	}

	gsl_vector_view fit_result_vec = gsl_vector_view_array(x, fd->nrows);

	gsl_multifit_linear(&(fd->lin_fit_params).matrix, &fit_result_vec.vector,fd->lin_fit_result, fd->cov, &chisq, fd->workspace);
	gsl_blas_dgemv(CblasNoTrans, 1.0, &(fd->lin_fit_params).matrix, fd->lin_fit_result, -1.0, &fit_result_vec.vector);

	//result is now stored in x, as required by levmar fitting function
}

/******************************************************************************/

static PyObject * calcShiftSqueeze(PyObject *self, PyObject *args){

	PyObject *spec, *wavelengths, *fitting_wavelengths, *A_arr;
	fit_data_t *fd;
	double p[2]={0.0, 1.0};  //shift and squeeze initial guess
	double *fit_result;
	double opts[LM_OPTS_SZ], info[LM_INFO_SZ];
	int min_fit_idx=0;

	//parse the arguments passed to the function by Python - no increase in ref count
	if(!PyArg_ParseTuple(args, "OOOO", &spec, &wavelengths, &A_arr, &fitting_wavelengths)){
		PyErr_SetString(PyExc_ValueError,"Invalid parameters. Expecting 4 args of type numpy.array");
		return NULL;
	}

	//Check that the numpy arrays we have been passed are c-style, contiguous
	//and aligned
	if(!PyArray_ISCARRAY(A_arr)){
		PyErr_SetString(PyExc_ValueError,"Fitting values arrays (A) is not C compatible");
		return NULL;
	}
	if(!PyArray_ISCARRAY(spec)){
			PyErr_SetString(PyExc_ValueError,"Spectrum array is not C compatible");
			return NULL;
	}
	if(!PyArray_ISCARRAY(wavelengths)){
			PyErr_SetString(PyExc_ValueError,"Wavelengths array is not C compatible");
			return NULL;
	}
	if(!PyArray_ISCARRAY(fitting_wavelengths)){
			PyErr_SetString(PyExc_ValueError,"Fitting wavelengths array is not C compatible");
			return NULL;
	}

	//TODO - check that the arrays have the correct dimensions

	//construct the fitting data object
	fd = createFitDataObj(A_arr, spec, wavelengths, fitting_wavelengths);

	//find the index of the minimum fitting wavelength in the wavelengths array
	while (fd->wavelengths[min_fit_idx] < fd->fitting_wavelengths[0]){
		min_fit_idx ++;
	}

	//printf("min_fit_idx = %d\n", min_fit_idx);
	//printf("min wavelength = %lf\n",fd->wavelengths[ min_fit_idx]);
//
//	//TODO - check that this is consistent with size of A and that idx is not 0
//
	//create array to hold fit result
	if ((fit_result = calloc(fd->nrows,sizeof(double)))==NULL){
		PyErr_SetString(PyExc_MemoryError,"Failed to allocate memory for fit result");
		return NULL;
	}


	//specify the fitting parameters for the Levenberg-Marquardt optimisation
	//TODO - comment what these are
	opts[0]=0.1;//LM_INIT_MU;
	opts[1]=1E-10;
	opts[2]=1E-10;
	opts[3]=1E-10;
	opts[4]= 1.0;//LM_DIFF_DELTA;

	//printf("LM_INIT_MU = %lf\n", LM_INIT_MU);
	//printf("LM_DIFF_DELTA = %lf\n", LM_DIFF_DELTA);

	//specify limits for optimisation parameters
	double lb[2], ub[2];
	lb[0]=-1.5; lb[1]=1.0;
	ub[0]=1.5;
	ub[1]=1.0;

	//Do the Levenberg-Marquardt optimisation to solve for shift and squeeze
	dlevmar_dif(fit_func, p, NULL, 1, fd->nrows, 1000, opts, info, NULL, NULL, fd);
	//printf("finished because: %lf\n", info[6]);
	//printf("shift = %lf, squeeze = %lf", p[0], p[1]);

	//tidy up
	freeFitData(&fd);
	free(fit_result);

	//create a Python tuple of shift and squeeze values and return it
	return Py_BuildValue("dd",p[0],p[1]);
}
/******************************************************************************/
static PyObject * getFitAt(PyObject *self, PyObject *args){

	PyObject *spec, *wavelengths, *fitting_wavelengths, *A_arr;
	fit_data_t *fd;
	double p[2];  //shift and squeeze initial guess
	double *fit_result;
	int min_fit_idx=0;

	//parse the arguments passed to the function by Python - no increase in ref count
	if(!PyArg_ParseTuple(args, "OOOOdd", &spec, &wavelengths, &A_arr, &fitting_wavelengths,&(p[0]),&(p[1]))){
		PyErr_SetString(PyExc_ValueError,"Invalid parameters. Expecting 4 args of type numpy.array");
		return NULL;
	}

	//Check that the numpy arrays we have been passed are c-style, contiguous
	//and aligned
	if(!PyArray_ISCARRAY(A_arr)){
		PyErr_SetString(PyExc_ValueError,"Fitting values arrays (A) is not C compatible");
		return NULL;
	}
	if(!PyArray_ISCARRAY(spec)){
			PyErr_SetString(PyExc_ValueError,"Spectrum array is not C compatible");
			return NULL;
	}
	if(!PyArray_ISCARRAY(wavelengths)){
			PyErr_SetString(PyExc_ValueError,"Wavelengths array is not C compatible");
			return NULL;
	}
	if(!PyArray_ISCARRAY(fitting_wavelengths)){
			PyErr_SetString(PyExc_ValueError,"Fitting wavelengths array is not C compatible");
			return NULL;
	}

	//TODO - check that the arrays have the correct dimensions

	//construct the fitting data object
	fd = createFitDataObj(A_arr, spec, wavelengths, fitting_wavelengths);

	//find the index of the minimum fitting wavelength in the wavelengths array
	while (fd->wavelengths[min_fit_idx] < fd->fitting_wavelengths[0]){
		min_fit_idx ++;
	}

	printf("min_fit_idx = %d\n", min_fit_idx);
	printf("min wavelength = %lf\n",fd->wavelengths[min_fit_idx]);
//
//	//TODO - check that this is consistent with size of A and that idx is not 0
//
	//create array to hold fit result
	if ((fit_result = malloc(sizeof(double) * fd->nrows))==NULL){
		PyErr_SetString(PyExc_MemoryError,"Failed to allocate memory for fit result");
		return NULL;
	}

	//initialise the fit_result with the measured spectrum linear interpolated
	//onto the fitting wavelength points
	if (linearInterpolate(fd->wavelengths, fd->measured_spec, fd->spec_size,
			              fd->fitting_wavelengths, fit_result, fd->nrows)){
		PyErr_SetString(PyExc_ValueError,"Failed to interpolate the spectrum onto the fitting points.");
		if(fit_result != NULL){
			free(fit_result);
		}
		return NULL;
	}

	//create numpy array of initial spec values
	npy_intp dims[2];
	dims[0] = (npy_intp)(fd->nrows);
	double *tmp_data = malloc(fd->nrows*sizeof(double));
	double *tmp_data2 = malloc(fd->nrows*sizeof(double));
	memcpy(tmp_data,fit_result, fd->nrows*sizeof(double));
	PyObject *raw_values, *fit_values;

	raw_values = createAutoDeallocNumpyArray(1, &(dims[0]), NPY_FLOAT64, tmp_data);

	fit_func(p, fit_result, 2, fd->nrows, fd);
	memcpy(tmp_data2,fit_result, fd->nrows*sizeof(double));
	fit_values = createAutoDeallocNumpyArray(1, &(dims[0]), NPY_FLOAT64, tmp_data2);

	//tidy up - note that tmp_data and tmp_data2 will be free'd by Python
	freeFitData(&fd);
	free(fit_result);

	//create a Python tuple of shift and squeeze values and return it
	return Py_BuildValue("OO",raw_values,fit_values);
}
/******************************************************************************/






//set up the functions to be visible in Python
static PyMethodDef shift_and_squeeze_methods[] = {
    {"calc_sns", calcShiftSqueeze, METH_VARARGS, ""},
    {"get_fit", getFitAt,METH_VARARGS, "" },
    {NULL, NULL}
};

//set up module to be importable in Python
PyMODINIT_FUNC initshift_and_squeeze(void){
    import_array(); //initialise NumPy API
    import_autodealloc_numpy();
    Py_InitModule3("shift_and_squeeze", shift_and_squeeze_methods,"");
}
